import cv2,os

def path_mkdir_if_notexists(path):
    if os.path.exists(path) == False:
        os.mkdir(path)


class VideoLoader:

    def __init__(self,video_path:str,img_cache:bool=True) -> None:
        self.img_cache = img_cache
        self.video_path = video_path
        self.cap = cv2.VideoCapture(self.video_path)
        # 视频总的帧数
        if img_cache:
            self.frame_count = self._cache()
        else:
            self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print(f"video all frame count is {self.frame_count}")
        pass

    def _cache(self):
        """
        缓冲所有的图片
        """
        # 图片缓冲池
        self.cache_pool = []
        frame_cnt = 0
        while True:
            ret,frame = self.cap.read()
            if ret == False:
                break
            self.cache_pool.append(frame)
            frame_cnt += 1
        # 记录缓冲的图片的个数
        return frame_cnt
    
    def getFrameImage(self,frameid_idx_from_0:int):
        """
        获取视频的特定帧,frameid index from 0.
        """
        frameid = frameid_idx_from_0
        if frameid < 0 or frameid >= self.frame_count:
            return False,None
        if self.img_cache:
            return True,self.cache_pool[frameid]
        # 下面这种调用方法存在坑
        if frameid != self.cap.get(cv2.CAP_PROP_POS_FRAMES):
            self.cap.set(cv2.CAP_PROP_POS_FRAMES,frameid)
        ret,img = self.cap.read()
        return ret,img



class VideoKeeper:

    def __init__(self,file_path:str,img_shape:tuple=None,fps:int=10) -> None:
        self.frame_id = 0
        self.file_path = file_path
        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.fps = fps
        self.video_size = img_shape
        if self.video_size is None:
            self.video = None
        else:
            self.video = cv2.VideoWriter(self.file_path,self.fourcc,self.fps,self.video_size,True)
        pass

    def _get_shape(self,img):
        self.video_size = ( img.shape[1] , img.shape[0] )
        print(f"shape -> {self.video_size}")

    def imwrite_from(self,img_path:str):
        # 从路径读取图片
        frame = cv2.imread(img_path)
        if self.video is None:
            self._get_shape(frame)
            self.video = cv2.VideoWriter(self.file_path,self.fourcc,self.fps,self.video_size,True)
        self.video.write(frame)
        self.frame_id += 1
    
    def imwrite(self,img):
        # 直接输出图片
        frame = img
        if self.video is None:
            self._get_shape(frame)
            self.video = cv2.VideoWriter(self.file_path,self.fourcc,self.fps,self.video_size,True)
        self.video.write(frame)
        self.frame_id += 1
    
    def save(self):
        self.video.release()
        print(f"all count frame is {self.frame_id}")
